module RbExt
  module Hash
    
    # Returns a copy of the hash containing only the key-value pairs
    # matched by the key names given to the method.
    #
    # Example:
    #   { :a => 1, :b => 2 }.only(:a)  # => { :a => 1 }
    #   { :abc => 1, :bbc => 2 }.only(/^ab/)    # => { :abc => 1 }
    #
    def only(*keys)
      if keys.first.is_a? Regexp
        reject { |key, value| key !~ keys.first }
      else
        reject { |key, value| !keys.include?(key) }
      end
    end
    
    # Returns a copy of the hash containing only the key-value pairs
    # that did not match the key names given to the method.
    #
    # Example:
    #   { :a => 1, :b => 2 }.except(:a)    # => { :b => 2 }
    #   { :abc => 1, :bbc => 2 }.except(/^ab/)   # => { :bbc => 2 }
    #
    def except(*keys)
      if keys.first.is_a? Regexp
        reject { |key, value| key =~ keys.first }
      else
        reject { |key, value| keys.include?(key) }
      end
    end
    
    # Returns an array containing the values of the hash
    # referenced by the keys given to the method.
    #
    # Example:
    #   { :a => 1, :b => 2, :c => 3 }.pick(:a, :c)    # => [ 1, 3 ]
    #
    def pick(*keys)
      return *keys.map { |key| self[key] } if keys.size > 1
      self[keys.first]
    end

    # Same return value as #pick but deletes the
    # the keys from the hash also.
    #
    def pick!(*keys)
      return *keys.map { |key| delete(key) } if keys.size > 1
      delete(keys.first)
    end
  end
end